package com.ayu.shortlink.engine.util;

import com.ayu.shortlink.common.cache.CaffeineCacheFactory;
import com.ayu.shortlink.common.cache.CommonCacheService;
import com.ayu.shortlink.engine.common.convention.exception.ServiceException;
import com.ayu.shortlink.engine.common.enums.BucketTokenKeyEnum;
import jodd.util.concurrent.ThreadFactoryBuilder;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Component;

import java.util.concurrent.*;

/**
 * 限流器工具类组件
 */
@Component
@Slf4j
@RequiredArgsConstructor
public class RateLimiterUtil implements CommandLineRunner {
    // 单次增加的token
    @Value("${rate-limiter.refill-tokens:5000}")
    private Integer refillTokens;
    // 增加的频率(2毫秒)
    @Value("${rate-limiter.refill-interval:2}")
    private Integer refillInterval;
    // 最大令牌容量
    @Value("${rate-limiter.bucket-capacity:15000}")
    private Integer bucketCapacity;

    private final CommonCacheService<String,Integer> tokenBucketCache = CaffeineCacheFactory.TOKEN_BUCKET_CACHE;

    private final ConcurrentMap<String,Object> LOCK_MAP = new ConcurrentHashMap<>();

    /**
     * 组件启动调用refillTokens()定时更新令牌
     * @param args args
     * @throws Exception 异常
     */
    @Override
    public void run(String... args) throws Exception {
        // 进行LOCK的初始化
        for (BucketTokenKeyEnum x : BucketTokenKeyEnum.values()){
            Object object = new Object();
            LOCK_MAP.put(x.getKey(), object);
        }
        // 开启定时任务
        refillTokens();
    }

    /**
     * 尝试获取一个令牌
     * @param key key
     * @return true - 获取成功 false - 获取失败
     */
    public boolean tryAcquire(String key){
        if (!LOCK_MAP.containsKey(key)){
            throw new ServiceException("当前key设置错误");
        }

        Integer result = tokenBucketCache.getMap().compute(key, (k, v) -> {
            if (v == null || v < 1) {
                return -1;
            }
            return v - 1;
        });
        return result >= 0;
    }

    /**
     * 尝试获取n个令牌
     * @param key key
     * @param tokenNumber 令牌数量
     * @return true - 获取成功 false - 获取失败
     */
    public boolean tryAcquire(String key,Integer tokenNumber){
        if (!LOCK_MAP.containsKey(key)){
            throw new ServiceException("当前key设置错误");
        }

        Integer result = tokenBucketCache.getMap().compute(key, (k, v) -> {
            if (v == null || v < tokenNumber) {
                return -1;
            }
            return v - tokenNumber;
        });
        return result >= 0;
    }

    /**
     * 在给定时间内尝试获取1个令牌
     * @param key key
     * @param timeOut 时间
     * @param timeUnit 时间单位
     * @return true - 获取成功 false - 获取失败
     */
    @SneakyThrows
    public boolean tryAcquire(String key, Integer timeOut, TimeUnit timeUnit){
        if (timeOut < 0 || timeUnit == null){
            throw new ServiceException("参数设置异常");
        }
        if (!LOCK_MAP.containsKey(key)){
            throw new ServiceException("当前key设置错误");
        }

        // 转换为nanos
        long nanos = timeUnit.toNanos(timeOut);
        long start = System.currentTimeMillis();
        long end = start + nanos;

        boolean result = false;
        // 获取当前key的锁
        Object LOCK = LOCK_MAP.get(key);
        // 在给定时间内
        while (System.currentTimeMillis() <= end){
            synchronized (LOCK){
                // 尝试扣减token
                Integer currentToken = tokenBucketCache.get(key);
                if (currentToken >= 0){
                    // 令牌数量够 可以执行扣减
                    tokenBucketCache.put(key, currentToken-1);
                    result = true;
                    break;
                }
                // 数量不够 进行等待
                long waitTime = end - System.currentTimeMillis();
                if (waitTime > 0){
                    LOCK.wait(waitTime);
                }
            }
        }
        return result;
    }

    /**
     * 在给定时间内尝试获取tokenNumber个令牌
     * @param key key
     * @param tokenNumber 令牌数量
     * @param timeOut 时间
     * @param timeUnit 时间单位
     * @return true - 获取成功 false - 获取失败
     */
    @SneakyThrows
    public boolean tryAcquire(String key,Integer tokenNumber,Integer timeOut, TimeUnit timeUnit){
        if (timeOut < 0 || timeUnit == null){
            throw new ServiceException("参数设置异常");
        }
        if (!LOCK_MAP.containsKey(key)){
            throw new ServiceException("当前key设置错误");
        }

        // 转换为nanos
        long nanos = timeUnit.toNanos(timeOut);
        long start = System.currentTimeMillis();
        long end = start + nanos;

        boolean result = false;
        // 获取当前key的锁
        Object LOCK = LOCK_MAP.get(key);
        // 在给定时间内
        while (System.currentTimeMillis() <= end){
            synchronized (LOCK){
                // 尝试扣减token
                Integer currentToken = tokenBucketCache.get(key);
                if (currentToken >= tokenNumber){
                    // 令牌数量够 可以执行扣减
                    tokenBucketCache.put(key, currentToken-tokenNumber);
                    result = true;
                    break;
                }
                // 数量不够 进行等待
                long waitTime = end - System.currentTimeMillis();
                if (waitTime > 0){
                    LOCK.wait(waitTime);
                }
            }
        }
        return result;
    }

    /**
     * 定时向令牌桶添加令牌
     */
    private void refillTokens(){
        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(
                ThreadFactoryBuilder.create()
                        .setDaemon(true)
                        .setUncaughtExceptionHandler((t, e) -> {
                            // 报警
                            log.error("定时向令牌桶添加令牌出现异常");
                            e.printStackTrace();
                        })
                        .setNameFormat("thread-refillTokens")
                        .get()
        );

        scheduler.scheduleAtFixedRate(() -> {
            try {
                // 进行令牌添加操作
                for (BucketTokenKeyEnum x : BucketTokenKeyEnum.values()){
                    int oldTokens = tokenBucketCache.get(x.getKey()) == null ? 0 : tokenBucketCache.get(x.getKey());
                    int newTokens = Math.min(bucketCapacity, oldTokens + refillTokens);
                    tokenBucketCache.put(x.getKey(), newTokens);
                    // 添加令牌之后 进行唤醒
                    Object LOCK = LOCK_MAP.get(x.getKey());
                    synchronized (LOCK){
                        LOCK.notifyAll();
                    }
                }
            } catch (Exception e) {
                log.error("令牌填充任务执行异常", e);
            }
        }, 0, refillInterval, TimeUnit.MILLISECONDS);
    }
}
